from document_loader import DocumentLoader
from embedding_client import EmbeddingClient
from vector_store import VectorStore
from retriever import Retriever
from prompt_builder import PromptBuilder
from openai_client import OpenAIClient
from dotenv import load_dotenv
import os

def main():


    # 加载 .env 文件
    #load_dotenv()
    # =========================

    # 现在你可以像访问普通环境变量一样访问 .env 文件中的变量
    ai_key = os.getenv('AI_QW_KEY')
    ai_url = os.getenv('AI_QW_URL')
    java = os.getenv('JAVA_HOME')
    print(java)
    print(ai_key)
    print(ai_url)
    # 初始化模块
    pdf_path = "../data/quanliucun.pdf"

    document_loader = DocumentLoader(pdf_path)
    embedding_client = EmbeddingClient(ai_key,ai_url)
    vector_store = VectorStore("rag_collection",embedding_client)
    retriever = Retriever(embedding_client, vector_store)
    prompt_builder = PromptBuilder()
    openai_client = OpenAIClient(ai_key,ai_url)

    # 加载并处理文档
    texts = document_loader.load_and_split(page_numbers=[0,2])
    #print("texts:::::::::",texts)
    for txt in texts:
        #print([txt])
        vector_store.add_documents([txt])

    print(">>>>>>>>>>>>>>RAG数据准备就绪<<<<<<<<<<<<<<")
    print(">>>>>>>>>>>>>>进入问答系统<<<<<<<<<<<<<<")
    while True:
        user_input = input("请输入你的问题（输入exit退出）: ")
        if user_input.lower() == 'exit':  # 不区分大小写退出
            print("程序已退出")
            break
            # 用户交互
        relevant_docs = retriever.retrieve(user_input)
        context = "\n".join(relevant_docs)
        prompt = prompt_builder.build_prompt(query=user_input, context=context)
        answer = openai_client.get_completion(prompt)

        for chunk in answer:
            # print(chunk.choices[0].delta)
            print(chunk.choices[0].delta.content, end="")

        print()



if __name__ == "__main__":
    main()